#!/usr/bin/env python
#
# Simulated grid (top-level component)
#
# Author:   Mike Murphy <mamurph@cs.clemson.edu>
# Revision: 23 September 2009
#
#   Copyright 2009 Clemson University
#
#   Licensed under the Apache License, Version 2.0 (the "License");
#   you may not use this file except in compliance with the License.
#   You may obtain a copy of the License at
#
#       http://www.apache.org/licenses/LICENSE-2.0
#
#   Unless required by applicable law or agreed to in writing, software
#   distributed under the License is distributed on an "AS IS" BASIS,
#   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#   See the License for the specific language governing permissions and
#   limitations under the License.
#
# This material is based upon work supported under a National Science
# Foundation Graduate Research Fellowship. Any opinions, findings,
# conclusions or recommendations expressed in this publication are those of
# the author(s) and do not necessarily reflect the views of the National
# Science Foundation.

'''
This module contains the top-level "Grid" class that links together Sites
and Clusters to form a simulated computational grid. Each grid consists of
a set of Sites, which represent Computing Elements (CEs) on physical grids.
Each Site is a named interface to a single scheduler on a cluster. Multiple
Sites may share the same cluster.
'''

from ce_sim import Site
from cluster_sim import Cluster, MachineFactory
from results import JobResult
from storage import FileCache


class Grid(object):
   '''
   Top-level grid simulation component: handles dynamic changes to the grid
   configuration as well as high-level job routing.
   '''
   def __init__(self, name, simobj, sched_trans, metasched=None, \
                machine_factory=MachineFactory(), ce_speed=-1, **kwds):
      '''
      @param name: Name of the grid system
      @type name: str
      @param simobj: Reference to the simulator object
      @param sched_trans: Scheduler translation function (fn pointer), which
                          instantiates a scheduler based on a single string
                          argument containing the name of the scheduler to
                          create
      @param metasched: Top-level grid metascheduler, which will use every
                        grid site as a metascheduling target (optional)
      @param machine_factory: MachineFactory instance to use when creating new
                              clusters (see L{cluster_sim.MachineFactory})
      @param ce_speed: Default speed (bytes/sec) of file caches on new CEs
                       added to the grid with the add_ce method. Set to zero
                       zero for instantaneous transfers, -1 to inhibit
                       creation of FileCache objects.
      @type ce_speed: number
      @param kwds: Extra data to be stored in the Grid object
      '''
      self.name = name
      self.sitemap = {}
      self.clustermap = {}
      self.sim = simobj
      self.sched_trans = sched_trans
      self.metasched = metasched
      self.factory = machine_factory
      self.ce_speed = ce_speed
      self.extra_data = kwds
   def add_site(self, site):
      '''
      Adds an existing Site object to the Grid (and metascheduler, if used).
      If a new Site needs to be created, use add_ce instead.

      @param site: Site object to add to the Grid
      '''
      site_name = site.name
      if site_name in self.sitemap:
         raise Exception('Site names must be unique')
      else:
         self.sitemap[site_name] = site
         if self.metasched is not None:
            self.metasched.add_target(site)
   def remove_site(self, site):
      '''
      Removes a Site from the Grid (and metascheduler, if used) using a
      reference to the Site object. To remove a Site by name, use remove_ce
      instead.

      @param site: Site object to remove from the Grid
      '''
      site_name = site.name
      del self.sitemap[site_name]
      if self.metasched is not None:
         self.metasched.remove_target(site)
   def get_site(self, name):
      '''
      @param name: Name of the site to obtain
      @type name: str

      @return: reference to a Site object known to the grid, or raises an
      Exception if no such site exists
      '''
      return self.sitemap[name]
   def register_cluster(self, cluster):
      '''
      Registers a Cluster object with the cluster tracking map. To create a
      new Cluster object, simply call resize_cluster with a cluster_name not
      currently tracked by the Grid.

      @param cluster: Cluster object to track
      '''
      name = cluster.name
      if name not in self.clustermap:
         self.clustermap[name] = cluster
   def get_waiting_count(self):
      '''
      Returns total number of jobs waiting across all sites
      '''
      count = 0
      for site_name in self.sitemap:
         count += self.sitemap[site_name].get_waiting_count()
      return count
   def get_total_queue_size(self):
      '''
      Returns total queue size (waiting+running) across all sites
      '''
      count = 0
      for site_name in self.sitemap:
         count += self.sitemap[site_name].get_total_queue_size()
      return count
   def get_completed_count(self):
      '''
      Returns total number of jobs completed on all sites
      '''
      count = 0
      for site_name in self.sitemap:
         count += self.sitemap[site_name].get_completed_count()
      return count
   def submit_job(self, job, target):
      '''
      Submits a job to a target Site, if the Site name matches a Site
      currently known to the Grid. If the name match fails, the job is
      flagged with an error condition, and a JobResult is created and sent
      to the simulator so that the error can be recorded.

      @param job: Job object to submit
      @param target: Site name to which the job should be sent. If the site
                     name is an empty string ("") and the optional grid-wide
                     metascheduler is in use, then the job will be submitted
                     to the metascheduler.
      @type target: str

      @return: simulator results collection
      '''
      r = []
      job.extra_data['grid_job'] = True
      job.extra_data['grid_name'] = self.name
      #
      if (target == '') and (self.metasched is not None):
         r.append(self.metasched.submit_job(job))
      elif target not in self.sitemap:
         # unmatchable job
         job.error = True
         r.append(JobResult(self.sim.get_time(), job))
      else:
         r.append(self.sitemap[target].submit_job(job))
      return (r,)
   def add_submit_event(self, time, job, target=''):
      '''
      Schedules a job submission event at a time not in the past.

      @param time: Time at which job submission should occur
      @type time: number
      @param job: Job object to submit
      @param target: Name of the site to which the job should be sent, or
                     grid-wide metascheduler if empty string
      @type target: str
      '''
      self.sim.new_external_event(time, self.submit_job, job, target)
   #
   # Dynamic grid methods
   def add_ce(self, ce_name, cluster_name, sched_name, speed=None):
      '''
      Creates a new Site and adds the Site to the Grid. If the Cluster
      providing machines to the Site does not have the specified scheduler,
      then a new scheduler is also created.

      @param ce_name: Name of the new Site
      @type ce_name: str
      @param cluster_name: Name of the Cluster that provides machines to the
                           Site. This Cluster MUST be present in the cluster
                           map, or an Exception will be raised.
      @type cluster_name: str
      @param sched_name: Name of the scheduler to be used with the Site. If
                         the scheduler is not already present on the back-end
                         Cluster, it will be created using the scheduler
                         translation function specified for this Grid.
      @type sched_name: str
      @param speed: Speed of site file cache, which will be added to the
                    scheduler's resource broker as the "vmstore" resource.
                    Set to zero for instantaneous transfers, or a negative
                    number to disable cache creation. Set to None (default)
                    to use the default speed set in the grid constructor.
      @type speed: number

      @return: reference to the newly created site
      '''
      cluster = self.clustermap[cluster_name]
      sched = None
      if sched_name not in cluster.schedulers:
         sched = self.sched_trans(sched_name, ce_name, self.sim)
         cluster.add_scheduler(sched_name, sched)
         sched.add_targets(cluster.machines)
      else:
         sched = cluster.schedulers[sched_name]
      s = Site(ce_name, sched, cluster)
      fspeed = self.ce_speed
      if speed is not None:
         fspeed = speed
      if fspeed >= 0:
         fc = FileCache('VMCache', self.sim, fspeed)
         sched.add_resource('vmstore', fc)
      self.add_site(s)
      return s
   def remove_ce(self, ce_name):
      '''
      Remove a Site from the Grid. Upon removal, any jobs running on the Site
      will be terminated and given an error status.

      @param ce_name: Name of the Site to remove
      @type ce_name: str

      @return: reference to the deleted site
      '''
      s = self.sitemap[ce_name]
      s.terminate()
      self.remove_site(s)
      return s
   def resize_cluster(self, cluster_name, cpu_count, profile=''):
      '''
      Resize an existing Cluster or create, register, and size a new Cluster,
      if the cluster name is not already known to the Grid. The resulting
      Cluster size must not be negative, or an Exception will be raised. When
      removing CPU cores from a cluster, idle cores will be selected for
      removal, if available. If there are insufficient idle cores to satisfy
      the removal request, then cores will be selected at random, and the jobs
      running on those cores will be killed.

      @param cluster_name: Name of Cluster to resize. If this Cluster does
                           not exist, it will be created with zero CPU cores
                           and then immediately resized.
      @type cluster_name: str
      @param cpu_count: Number of CPU cores to add (positive) or remove
                        (negative) from the Cluster.
      @type cpu_count: int
      @param profile: Profile name for any new machines
      @type profile: str

      @return: reference to the cluster
      '''
      if cluster_name not in self.clustermap:
         cluster = Cluster(cluster_name, self.sim, factory=self.factory)
         self.register_cluster(cluster)
      if cpu_count > 0:
         machines = self.clustermap[cluster_name].grow(cpu_count, profile)
         for m in machines:
            for s in self.clustermap[cluster_name].schedulers:
               self.clustermap[cluster_name].schedulers[s].add_target(m)
      elif cpu_count < 0:
         machines = self.clustermap[cluster_name].shrink(cpu_count)
         for m in machines:
            for s in self.clustermap[cluster_name].schedulers:
               self.clustermap[cluster_name].schedulers[s].remove_target(m)
      return self.clustermap[cluster_name]
   #
#


if __name__ == '__main__':
   from egee_schedulers import new_egee_scheduler
   from job_model import Job
   from kernel import Simulator
   from metascheduler import Metascheduler
   from storage import FileStore
   from units import *
   from virtualization import VMHost
   from vo_format_reader import VO
   from voc_sim import AutoVOC
   from watchdog_sim import Watchdog

   sim = Simulator()
   mf = MachineFactory(VMHost, vmm_name='KVM', sim=sim, cpu_oh=9.0)
   ms = Metascheduler(sim)
   fs = FileStore('VMStore', sim, 1*u_Gbps)

   grid = Grid('Grid', sim, new_egee_scheduler, metasched=ms, \
               machine_factory=mf, ce_speed=1*u_Gbps)
   grid.resize_cluster('cluster1', 10)
   grid.add_ce('site1', 'cluster1', 'condor')

   engage = VO('engage')
   av = AutoVOC(grid, fs, 5*u_GiB, 'condor', 60, Watchdog, sim=sim, \
                q_sched=grid.get_site('site1').sched)
   av['engage'] = engage

   for x in xrange(0,20):
      job = Job(20 + (x % 4), name='job%d' % x, vo='engage', prob_err=0.2)
      grid.add_submit_event(10*x, job, 'voc_ce:engage')
   #

   sim.run()
#
